import os,sys,math,time,linecache
from math import ceil
from operator import mod
from pathlib import Path
import cv2
from PIL import Image
import numpy as np
#import PySimpleGUI as sg
from matplotlib import pyplot as plt
from matplotlib import pylab as pylab
from matplotlib.lines import Line2D as line
from matplotlib.patches import Rectangle as box
import linecache

try:
    media_file = sys.argv[1]
except:
    media_file = 'test.jpg'

classNames = ['Motor','Mount']
colors = ['green','red']
labelFile = ''.join([media_file,'.label'])

def get_labels(media_file):
    labelRows = open(labelFile).read().strip().split('\n')
    labelRows = linecache.getlines(labelFile)
    classes = [r[r.find(' ') + 1:].split(',')[0] for r in labelRows]

    return classes

def put_label(media_file,left_top,right_bottom,className=None):
    labelFile = ''.join([media_file,'.label'])
    with open(labelFile,'a') as f:
        if className==None:
            f.writelines(''.join(['className',':\t',str(left_top),str(right_bottom),'\n']))
        else:
            f.writelines(''.join([className,':\t',str(left_top),str(right_bottom),'\n']))
    
    return 

def draw_box(image,left,top,right,bottom,color):
    width,height = image.size
    left_bottom = [min(left,right),min(top,bottom)]
    width_box = abs(right - left)
    height_box = abs(top - bottom)
    if color == 'red':
        color = 'r'
    elif color == 'blue':
        color = 'b'
    elif color == 'green':
        color = 'g'
    elif color == 'yellow':
        color = 'y'
    alpha = 0.4
    rect = box(left_bottom,width_box,height_box,color=color,alpha=alpha)
    ax.add_patch(rect)

    return rect
    
def remove_box(box_obj):
    try:
        ax.patches.remove(box_obj)
    except:
        pass

    return 

def draw_tag(img,tag,position_row_tag,position_col_tag):
    try:
        if tag == classNames[0]:
            tag_color = 'g'
        else:
            tag_color = 'r'

        tag = ax.text(position_row_tag,position_col_tag,tag,color=tag_color)
    except:
        print("draw tag failed.")
    return tag

def remove_tag(tag_obj):
    try:
        ax.texts.remove(tag_obj)
    except:
        print(dir(ax))
        pass

    return 
def load_img_by_cv(media_file):
    global image_0
    global image_1
    global image_2
    global img_0
    global img_1
    global img_2
    image_0 = cv2.imread(media_file,cv2.IMREAD_GRAYSCALE)
    image_1 = cv2.imread(media_file,cv2.IMREAD_COLOR)
    image_2 = cv2.imread(media_file,cv2.IMREAD_UNCHANGED)
    img_0 = image_0[...,::-1]
    img_1 = image_1[...,::-1]
    img_2 = image_2[...,::-1]

def load_img_by_pillow(media_file):
    image = Image.open(media_file)

    return image

def draw_line(image,position,ROW_or_COL,REF_or_TAG):
    width,height = image.size
    if ROW_or_COL == 1:
        LINE = [(position,1),(position,height)]
    else:
        LINE = [(1,position),(width,position)]
    (line_xs,line_ys) = zip(*LINE)
    if REF_or_TAG == 0:
        color='red'
        sort='dashed'
        linewidth=0.5
    else:
        color='blue'
        sort='solid'
        linewidth=0.5
    line_obj = line(line_xs,line_ys,linewidth,sort,color)
    ax.add_line(line_obj)
    return line_obj 

def remove_line(line_obj):
    try:
        ax.lines.remove(line_obj)
    except:
        pass
    return 
    
def draw_REF_lines(image):
    width,height = image.size
    deci_width = int(ceil(width/10))
    deci_height = int(ceil(height/10))
    REF_width_positions = [
        deci_width*1,deci_width*2,deci_width*3,deci_width*4,deci_width*5,
        deci_width*6,deci_width*7,deci_width*8,deci_width*9
        ]
    REF_height_positions = [
        deci_height*1,deci_height*2,deci_height*3,deci_height*4,deci_height*5,
        deci_height*6,deci_height*7,deci_height*8,deci_height*9
        ]
    for position in REF_width_positions:
        draw_line(image,position,COL,REF)
        #REF_tag_row = ax.text(position,position_col,tag,color=tag_color)
    for position in REF_height_positions:
        draw_line(image,position,ROW,REF)
        REF_tag_col = ax.text(10,position - deci_height/2 + 11 ,REF_height_positions.index(position)+1,color='r')
    REF_tag_col = ax.text(10,deci_height*10 - deci_height/2 + 11,0,color='r')
    for index in range(0,10):
        if index == 0:
            tag = 'Q'
        elif index == 1:
            tag = 'W'
        elif index == 2:
            tag = 'E'
        elif index == 3:
            tag = 'R'
        elif index == 4:
            tag = 'T'
        elif index == 5:
            tag = 'Y'
        elif index == 6:
            tag = 'U'
        elif index == 7:
            tag = 'I'
        elif index == 8:
            tag = 'O'
        elif index == 9:
            tag = 'P'
        REF_tag_row = ax.text(index*deci_width + deci_width/2 -11 ,20,tag,color='r')

    return

def get_position_by_num(image,num):
    width,height = image.size
    deci_width = int(ceil(width/10))
    deci_height = int(ceil(height/10))
    REF_width_positions = [
        deci_width*1,deci_width*2,deci_width*3,deci_width*4,deci_width*5,
        deci_width*6,deci_width*7,deci_width*8,deci_width*9,deci_width*10,
        ]
    REF_height_positions = [
        deci_height*1,deci_height*2,deci_height*3,deci_height*4,deci_height*5,
        deci_height*6,deci_height*7,deci_height*8,deci_height*9,deci_height*10,
        ]
    if num == 1:
        position_width = 0 + ceil(REF_width_positions[0]/2)
        position_height = 0 + ceil(REF_height_positions[0]/2)
    elif num == 0:
        position_width = REF_width_positions[10-2] + ceil((REF_width_positions[10-1]-REF_width_positions[10-2])/2)
        position_height = REF_height_positions[10-2] + ceil((REF_height_positions[10-1]-REF_height_positions[10-2])/2)
    else:
        position_width = REF_width_positions[num-2] + ceil((REF_width_positions[num-1]-REF_width_positions[num-2])/2)
        position_height = REF_height_positions[num-2] + ceil((REF_height_positions[num-1]-REF_height_positions[num-2])/2)
    return position_width,position_height

lines_COL = []
lines_ROW = []
boxes_RED = []
boxes_GREEN = []
tagsFirst = []
tagsSecond = []
startOrEnd = 0
className = classNames[0]
color = colors[0]
positionCol = 0
positionRow = 0
IMAGE = load_img_by_pillow(media_file)
width,height = IMAGE.size
deci_width = int(ceil(width/10))
deci_height = int(ceil(height/10))

def on_key_press(event):
    num = 10
    global positionCol
    global positionRow
    global startOrEnd
    global className
    global color
    global lines_ROW
    global lines_COL
    global boxes_RED
    global boxes_GREEN
    global tagsFirst
    global tagsSecond
    #global deci_width
    #global deci_height

    if event.key in "1234567890":
        num = int(event.key)
        positionCol,positionRow = get_position_by_num(IMAGE,num)
        if len(lines_ROW) > 0:
            remove_line(lines_ROW[-1])
        new_line = draw_line(IMAGE,positionRow,ROW,TAG)
        lines_ROW.append(new_line)
    elif event.key in "qwertyuiop":
        if event.key == 'q':
            num = 1
        if event.key == 'w':
            num = 2
        if event.key == 'e':
            num = 3
        if event.key == 'r':
            num = 4
        if event.key == 't':
            num = 5
        if event.key == 'y':
            num = 6
        if event.key == 'u':
            num = 7
        if event.key == 'i':
            num = 8
        if event.key == 'o':
            num = 9
        if event.key == 'p':
            num = 0
        positionCol,positionRol = get_position_by_num(IMAGE,num)
        if len(lines_COL) > 0:
            remove_line(lines_COL[-1])
        new_line = draw_line(IMAGE,positionCol,COL,TAG)
        lines_COL.append(new_line)
    elif event.key in "asdfghjkl;":
        if event.key == 'a':
            positionCol -= ceil(deci_width/2) 
            if len(lines_COL) > 0:
                for line in lines_COL:
                    remove_line(line)
            moved_line_COL = draw_line(IMAGE,positionCol,COL,TAG)
            lines_COL.append(moved_line_COL)
        if event.key == 's':
            positionCol -= 15
            if len(lines_COL) > 0:
                for line in lines_COL:
                    remove_line(line)
            moved_line_COL = draw_line(IMAGE,positionCol,COL,TAG)
            lines_COL.append(moved_line_COL)
        if event.key == 'd':
            positionCol -= 5
            if len(lines_COL) > 0:
                for line in lines_COL:
                    remove_line(line)
            moved_line_COL = draw_line(IMAGE,positionCol,COL,TAG)
            lines_COL.append(moved_line_COL)
        if event.key == 'f':
            positionCol += 5
            if len(lines_COL) > 0:
                for line in lines_COL:
                    remove_line(line)
            moved_line_COL = draw_line(IMAGE,positionCol,COL,TAG)
            lines_COL.append(moved_line_COL)
        if event.key == 'g':
            positionCol += 15
            if len(lines_COL) > 0:
                for line in lines_COL:
                    remove_line(line)
            moved_line_COL = draw_line(IMAGE,positionCol,COL,TAG)
            lines_COL.append(moved_line_COL)
        if event.key == 'h':
            positionRow += 15
            if len(lines_ROW) > 0:
                for line in lines_ROW:
                    remove_line(line)
            moved_line_ROW = draw_line(IMAGE,positionRow,ROW,TAG)
            lines_ROW.append(moved_line_ROW)
        if event.key == 'j':
            positionRow += 5
            if len(lines_ROW) > 0:
                for line in lines_ROW:
                    remove_line(line)
            moved_line_ROW = draw_line(IMAGE,positionRow,ROW,TAG)
            lines_ROW.append(moved_line_ROW)
        if event.key == 'k':
            positionRow -= 5
            if len(lines_ROW) > 0:
                for line in lines_ROW:
                    remove_line(line)
            moved_line_ROW = draw_line(IMAGE,positionRow,ROW,TAG)
            lines_ROW.append(moved_line_ROW)
        if event.key == 'l':
            positionRow -= 15
            if len(lines_ROW) > 0:
                for line in lines_ROW:
                    remove_line(line)
            moved_line_ROW = draw_line(IMAGE,positionRow,ROW,TAG)
            lines_ROW.append(moved_line_ROW)
        if event.key == ';':
            positionRow -= ceil(deci_height/2)
            if len(lines_ROW) > 0:
                for line in lines_ROW:
                    remove_line(line)
            moved_line_ROW = draw_line(IMAGE,positionRow,ROW,TAG)
            lines_ROW.append(moved_line_ROW)

    elif event.key in "[]":
        if event.key == '[':
            className = classNames[0]
            color = colors[0]
        elif event.key == ']':
            className = classNames[1]
            color = colors[1]
    elif event.key in "-+":
        if event.key == '-':
            pass
        elif event.key == '+':
            pass
    elif event.key == ' ':
        startOrEnd += 1
        REF_tag_row = ax.text(positionCol,positionRow,''.join(['[',str(positionCol),',',str(positionRow),']']),color='r')

        with open(labelFile,'a') as f:
            f.writelines(''.join([str(positionRow),',',str(positionCol),',']))
        if mod(startOrEnd,2) != 0:
            with open(labelFile,'a') as f:
                f.writelines(''.join([className,',']))
        if mod(startOrEnd,2) != 1:
            with open(labelFile,'a') as f:
                f.writelines(''.join(['\n']))
            with open(labelFile,'r') as f:
                labelInfo = f.read().splitlines()
            try:
                top = int(labelInfo[-1].split(',')[0])
                left= int(labelInfo[-1].split(',')[1])
                tag= labelInfo[-1].split(',')[2]
                bottom= int(labelInfo[-1].split(',')[3])
                right= int(labelInfo[-1].split(',')[4])
                newBox = draw_box(IMAGE,left,top,right,bottom,color)
                if color == colors[0]:
                    boxes_RED.append(newBox)
                else:
                    boxes_GREEN.append(newBox)
                #position_row = left
                position_row_tag = abs(ceil((left + right)/2)) - abs(ceil((right-left)/8))
                position_col_tag = ceil((top + bottom)/2)
                new_tag = draw_tag(IMAGE,className,position_row_tag,position_col_tag)
                if new_tag == classNames[0]:
                    tagsFirst.append(new_tag)
                else:
                    tagsSecond.append(new_tag)
            except:
                print("ERROR")

        positionRow,positionCol = 0,0
        if len(lines_ROW) > 0:
            for line in lines_ROW:
                remove_line(line)
        if len(lines_COL) > 0:
            for line in lines_COL:
                remove_line(line)
    elif event.key == '.':
        for index_line in lines_COL:
            remove_box(index_line)
        for index_line in lines_ROW:
            remove_box(index_line)
        for index_box in boxes_RED:
            remove_box(index_box)
        for index_box in boxes_GREEN:
            remove_box(index_box)
        for index_tag in tagsFirst:
            remove_tag(index_tag)
        for index_tag in tagsSecond:
            remove_tag(index_tag)
        positionRow,positionCol = 0,0

    elif event.key == '/':
        sys.exit()
    fig.canvas.draw_idle()

    return 

ROW,COL = 0,1
REF,TAG = 0,1
params = {
    'axes.labelsize':   '16',
    'xtick.labelsize':  '16',
    'ytick.labelsize':  '16',
    'lines.linewidth':  '2',
    'legend.fontsize':  '20',
    'figure.figsize':   '26,24',
    }

pylab.rcParams.update(params)
fig,ax = plt.subplots(1)


ax.imshow(IMAGE)
plt.axis('on')
plt.title(media_file)
draw_REF_lines(IMAGE)

fig.canvas.mpl_disconnect(fig.canvas.manager.key_press_handler_id)
fig.canvas.mpl_connect('key_press_event', on_key_press)
plt.show()
